#include "../rpc.pb.h"
#include "../echo_service.pb.h"
#include "../RpcChannel.h"
#include "../../TcpConnection.h"
#include "../../TcpClient.h"
#include "../../TcpServer.h"
#include "../../EventLoop.h"
#include <memory>
#include "../RpcController.h"
#include "../../Channel.h"
#include "../../noncopyable.h"

// #include "MurmurHash3.h"
#include <stdio.h>
#include "ConsistentHash.h"
#include "loadbalancer.h"

using namespace std::placeholders;

extern std::map<std::thread::id, std::shared_ptr<std::mutex>> readLockMap_;

class Balancer
{
public:
    Balancer(EventLoop* loop, 
            const string& ip, 
            const int& port, 
            const std::vector<std::pair<string, int>>& backends)
        : loop_(loop)
        , server_(loop, port)
        , backends_(backends)
        , current_(0)
        , consistentHash_(3)
    {
        init();
        server_.setConnectionCallback(std::bind(&Balancer::onConnection, this, _1));
        server_.setMessageCallback(std::bind(&Balancer::onMessage, this, _1, _2));
    }

    void setThreadNum(int numThreads) { server_.setThreadNum(numThreads); }
    void start() { server_.start(); }

private:

    void init()
    {

        for (size_t i = 0; i < backends_.size(); ++i)
        {
            backends_vec_.push_back( make_shared<BackendSession>(loop_, backends_[i].first, backends_[i].second, &consistentHash_));
            backends_vec_.back()->connect();
        }

    }

    void onConnection(const TcpConnectionPtr& conn)
    {
    }

    void onMessage(const TcpConnectionPtr& conn, Buffer* buf)
    {
        while (buf->readable()) {
            RpcMessagePtr rpcmessage( prototype_->New() );

            // std::string msg = buf->retrieveAllAsString();
            std::string msg;
            if (!parseMessage(conn, buf, msg)) {
                std::cout << "wait for more data\n";
                return ;
            }
            rpcmessage->ParseFromString(msg);
            // cout << to_string(rpcmessage->id()) + " message.service(): " + rpcmessage->service() + "\n";
            onRpcMessage(conn, rpcmessage);
        }
    }

    void onRpcMessage(const TcpConnectionPtr& conn, const RpcMessagePtr& msg)
    {
        std::lock_guard<std::mutex> lk(*readLockMap_[std::this_thread::get_id()]);
        // std::lock_guard<std::mutex> lk(globalMutex_);
        std::string key = conn->getPeerIpPort();
        consistentHash_.selectServer(key)->send(*msg, conn, key);

    }

    EventLoop* loop_;
    TcpServer server_;
    const RpcMessagePtr prototype_;

    std::vector<std::pair<string, int>> backends_;

    ConsistentHash consistentHash_;
    int current_;
    std::vector<std::shared_ptr<BackendSession>> backends_vec_;

    std::mutex globalMutex_;

};

int main(int argc, char* argv[]) 
{
    if (argc < 3)
    {
        fprintf(stderr, "Usage: %s listen_port backend_ip:port [backend_ip:port]\n", argv[0]);
        exit(0);
    }

    std::vector<std::pair<string, int>> backends;
    for (int i = 2; i < argc; ++i)
    {
        string ip = "127.0.0.1";
        int port = static_cast<int>(atoi(argv[i]));
        backends.push_back(std::pair(ip, port));

        std::cout << i << "th port : " << port << std::endl;

    }

    uint16_t port = static_cast<uint16_t>(atoi(argv[1]));
    // string ip = "127.0.0.1";
    std::cout << "listen port : " << port << std::endl;

    EventLoop loop;
    Balancer balancer(&loop, "127.0.0.1", port, backends);
    balancer.setThreadNum(4);
    // balancer.setThreadNum(0);
    balancer.start();
    loop.loop();


    return 0;
}

// g++ -std=c++17 *.cpp ../*.cpp ../../*.cpp -lpthread -lprotobuf -o balancer



